import asyncio
import json
import logging
from typing import List, Dict, Optional

import aiohttp
from pydantic import BaseModel

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class AgentTask(BaseModel):
    task_type: str
    parameters: Dict
    context: Optional[List[Dict]] = None


class AgentWorkflow(BaseModel):
    tasks: List[AgentTask]
    max_iterations: int = 3


class AgentClient:
    def __init__(self, base_url: str = "http://localhost:8000"):
        self.base_url = base_url
        self.session = None

    async def __aenter__(self):
        """创建异步会话"""
        if self.session is None:
            self.session = aiohttp.ClientSession()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """关闭异步会话"""
        if self.session:
            await self.session.close()
            self.session = None

    async def execute_workflow(self, workflow: AgentWorkflow) -> Dict:
        """执行智能体工作流"""
        if self.session is None:
            async with self as client:
                return await self._execute_workflow_internal(workflow)
        return await self._execute_workflow_internal(workflow)

    async def _execute_workflow_internal(self, workflow: AgentWorkflow) -> Dict:
        """内部工作流执行方法"""
        try:
            async with self.session.post(
                    f"{self.base_url}/v1/agent/workflow",
                    json=workflow.dict()
            ) as response:
                if response.status != 200:
                    error_text = await response.text()
                    raise Exception(f"工作流执行失败: {error_text}")
                return await response.json()
        except aiohttp.ClientError as e:
            logger.error(f"请求错误: {str(e)}")
            raise Exception(f"工作流执行失败: {str(e)}")
        except Exception as e:
            logger.error(f"执行错误: {str(e)}")
            raise

    async def execute_single_task(self, task: AgentTask) -> Dict:
        """执行单个任务"""
        workflow = AgentWorkflow(tasks=[task], max_iterations=1)
        result = await self.execute_workflow(workflow)
        return result["results"][0] if result["results"] else None


async def main():
    """主函数示例"""
    try:
        # 创建客户端实例
        async with AgentClient() as client:
            # 定义任务
            task1 = AgentTask(
                task_type="analysis",
                parameters={
                    "text": "分析以下文本的情感倾向：'今天天气真好，阳光明媚，心情愉快。'",
                    "aspects": ["情感", "强度"]
                }
            )

            task2 = AgentTask(
                task_type="summary",
                parameters={
                    "text": "这是一段需要总结的长文本...",
                    "max_length": 100
                }
            )

            # 创建工作流
            workflow = AgentWorkflow(
                tasks=[task1, task2],
                max_iterations=2
            )

            # 执行工作流
            logger.info("开始执行工作流...")
            results = await client.execute_workflow(workflow)
            logger.info("工作流执行完成")
            print("\n工作流执行结果:")
            print(json.dumps(results, ensure_ascii=False, indent=2))

            # 执行单个任务
            logger.info("开始执行单个任务...")
            single_result = await client.execute_single_task(task1)
            logger.info("单个任务执行完成")
            print("\n单个任务执行结果:")
            print(json.dumps(single_result, ensure_ascii=False, indent=2))

    except Exception as e:
        logger.error(f"执行出错: {str(e)}")
        raise


async def test_tools():
    """测试服务端工具调用接口"""
    async with aiohttp.ClientSession() as session:
        # 1. 测试 my_search_tool
        resp = await session.post(
            "http://localhost:8000/v1/agent/tool_call",
            json={"tool_name": "my_search_tool", "params": {"query": "python 是什么"}}
        )
        print("my_search_tool:", await resp.json())
        # 2. 测试 add
        resp = await session.post(
            "http://localhost:8000/v1/agent/tool_call",
            json={"tool_name": "add", "params": {"a": 3, "b": 5}}
        )
        print("add:", await resp.json())
        # 3. 测试 multiply
        resp = await session.post(
            "http://localhost:8000/v1/agent/tool_call",
            json={"tool_name": "multiply", "params": {"a": 4, "b": 6}}
        )
        print("multiply:", await resp.json())


if __name__ == "__main__":
    asyncio.run(main())
    asyncio.run(test_tools())
